%%% This script endogenizes lambda via innovation and knowledge diffusion

d_geo_const = exp(6.426134 / theta);
d_geo = d_geo_const*ones(N,N);
for i=1:N
   d_geo(i,i) = 1; % row: cited; column: citing
end

d_migr_const_aux = -(3.33587/theta);

d_geo_distance_const_aux = .0001645 / theta;
d_geo_distance = zeros(N,N);
for i=1:N
    for j=1:N
        d_geo_distance(i,j) = exp(d_geo_distance_const_aux*d_loc(i,j));
    end
end

figure
data_aux = load('dsets_for_model\macroipc_citing_cited_fe_cpp.csv');
d_know = reshape(exp(-data_aux(:,3)/theta),S,S); % NOTE: rows are origins (cited), columns are destinations (citing)
delta_r_to_s = reshape(-data_aux(:,3)/theta,S,S); 

XLabels = 1:S; YLabels = 1:S;
CustomXLabels = string(XLabels); CustomYLabels = string(YLabels);
CustomXLabels(1,1:11) = ["A1","A2","A3","B1","B2","C1","E1","F1","F2","G1","H1"];
CustomYLabels(1,1:11) = ["A1","A2","A3","B1","B2","C1","E1","F1","F2","G1","H1"];

heatmap(delta_r_to_s, 'Colormap', gray, 'XDisplayLabels', CustomXLabels, 'YDisplayLabels', CustomYLabels)
clear data_aux

ave_epsilon_aux_cz_dyn = zeros(N,T);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%% Estimate alpha and epsilon shocks
alpha_dyn = zeros(S,T); epsilon_dyn = zeros(N,S,T);

diffusion_frictions_dyn = zeros(N*S,N*S,T); % diffusion_frictions = (1/(d^G*d^M*d^K))^theta -> NOTE: rows are origins (cited), columns are destinations (citing)

ave_epsilon_aux_dyn = zeros(S,T);

%%% Start from t=first_period+1

for t=first_period:first_period+T_aux-1
    
    alpha_0 = ones(1,S); 
    
    d_migr = ones(N,N);
    for i=1:N
        for j=1:N
            d_migr(i,j) = exp(d_migr_const_aux*migration_exposure_dyn(i,j,t-1)); % The time period is t-1 because it applies to those who are old at t and hence migrated to j at t-1 when young
        end
    end
    
    for i=1:N
        for s=1:S 
            for j=1:N
                for r=1:S
                    diffusion_frictions_dyn((s-1)*N+i,(r-1)*N+j,t) = (1/(d_geo(i,j)*d_migr(i,j)*d_geo_distance(i,j)*d_know(s,r)))^theta;
                end
            end
        end
    end

    error_0 = 1000; count = 0;
    
    while error_0 > 0.0001
        count = count + 1;
        
        Summation = zeros(N,S);
        for j=1:N
            for r=1:S
                Summation(j,r) = sum(reshape(lambda_dyn(:,:,t-1),N*S,1).*...
                    diffusion_frictions_dyn(:,(r-1)*N+j,t).*reshape(repmat(alpha_0.^theta,N,1),N*S,1));
            end
        end
    
        % epsilon_aux = epsilon^theta
        % epsilon_aux_0 = max(0.1,(lambda_dyn(:,:,t) - lambda_dyn(:,:,t-1)) ./ Summation);
        epsilon_aux_0 = (lambda_dyn(:,:,t) - delta*lambda_dyn(:,:,t-1)) ./ Summation;
        
        ave_epsilon_aux = zeros(1,S);
        for s=1:S
            ave_epsilon_aux(s) = sum((epsilon_aux_0(:,s)).*L_y_dyn(:,s,t)) / sum(L_y_dyn(:,s,t));
        end

        alpha_1 = ave_epsilon_aux.*alpha_0;

        error_0 = max(abs(alpha_0 - alpha_1));

        alpha_0 = 0.8*alpha_0 + 0.2*alpha_1;

        if mod(count,10)==0
            error_0
        end

    end
    
    alpha_dyn(:,t) = alpha_0';
    epsilon_dyn(:,:,t) = epsilon_aux_0.^(1/theta);
    ave_epsilon_aux_dyn(:,t) = ave_epsilon_aux';
    
    for i=1:N
         ave_epsilon_aux_cz_dyn(i,t) = epsilon_aux_0(i,:)*(L_y_dyn(i,:,t)/sum(L_y_dyn(i,:,t)))';
    end
    
end

alpha_hat_dyn = log(alpha_dyn(:,first_period+1:first_period+T_aux-1)) - log(alpha_dyn(:,first_period:first_period+T_aux-2));

%%% Now, compute model with no knowledge flows across fields (intra):

diffusion_frictions_intra_dyn = zeros(N*S,N*S,T); 

alpha_intra_dyn = zeros(S,T); epsilon_intra_dyn = zeros(N,S,T);

for t=first_period:first_period+T_aux-1
    
    alpha_0 = ones(1,S); 
    
    d_migr = ones(N,N);
    for i=1:N
        for j=1:N
            d_migr(i,j) = exp(d_migr_const_aux*migration_exposure_dyn(i,j,t-1)); % The time period is t-1 because it applies to those who are old at t and hence migrated to j at t-1 when young
        end
    end
    
    for i=1:N
        for s=1:S 
            for j=1:N
                for r=1:S
                    if r == s
                        diffusion_frictions_intra_dyn((s-1)*N+i,(r-1)*N+j,t) = (1/(d_geo(i,j)*d_migr(i,j)*d_geo_distance(i,j)))^theta;
                    else
                        diffusion_frictions_intra_dyn((s-1)*N+i,(r-1)*N+j,t) = 0;
                    end
                end
            end
        end
    end

    error_0 = 1000; count = 0;
    
    while error_0 > 0.0001
        count = count + 1;
        
        Summation = zeros(N,S);
        for j=1:N
            for r=1:S
                Summation(j,r) = sum(reshape(lambda_dyn(:,:,t-1),N*S,1).*...
                    diffusion_frictions_intra_dyn(:,(r-1)*N+j,t).*reshape(repmat(alpha_0.^theta,N,1),N*S,1));
            end
        end
    
        % epsilon_aux = epsilon^theta
        epsilon_aux_0 = (lambda_dyn(:,:,t) - delta*lambda_dyn(:,:,t-1)) ./ Summation;
        
        ave_epsilon_aux = zeros(1,S);
        for s=1:S
            ave_epsilon_aux(s) = sum((epsilon_aux_0(:,s)).*L_y_dyn(:,s,t)) / sum(L_y_dyn(:,s,t));
        end

        alpha_1 = ave_epsilon_aux.*alpha_0;

        error_0 = max(abs(alpha_0 - alpha_1));

        alpha_0 = 0.8*alpha_0 + 0.2*alpha_1;

        if mod(count,10)==0
            error_0
        end

    end
    
    alpha_intra_dyn(:,t) = alpha_0';
    epsilon_intra_dyn(:,:,t) = epsilon_aux_0.^(1/theta);
    
end

alpha_hat_intra_dyn = log(alpha_intra_dyn(:,first_period+1:first_period+T_aux-1)) - log(alpha_intra_dyn(:,first_period:first_period+T_aux-2));

%%% Now, compute model with no knowledge flows due to migration (nomigr)

diffusion_frictions_nomigr_dyn = zeros(N*S,N*S,T); 

alpha_nomigr_dyn = zeros(S,T); epsilon_nomigr_dyn = zeros(N,S,T);

for t=first_period:first_period+T_aux-1
    
    alpha_0 = ones(1,S); 
    
    for i=1:N
        for s=1:S 
            for j=1:N
                for r=1:S
                    diffusion_frictions_nomigr_dyn((s-1)*N+i,(r-1)*N+j,t) = (1/(d_geo(i,j)*d_geo_distance(i,j)*d_know(s,r)))^theta;
                end
            end
        end
    end

    error_0 = 1000; count = 0;
    
    while error_0 > 0.0001
        count = count + 1;
        
        Summation = zeros(N,S);
        for j=1:N
            for r=1:S
                Summation(j,r) = sum(reshape(lambda_dyn(:,:,t-1),N*S,1).*...
                    diffusion_frictions_nomigr_dyn(:,(r-1)*N+j,t).*reshape(repmat(alpha_0.^theta,N,1),N*S,1));
            end
        end
    
        % epsilon_aux = epsilon^theta
        epsilon_aux_0 = (lambda_dyn(:,:,t) - delta*lambda_dyn(:,:,t-1)) ./ Summation;
        
        ave_epsilon_aux = zeros(1,S);
        for s=1:S
            ave_epsilon_aux(s) = sum((epsilon_aux_0(:,s)).*L_y_dyn(:,s,t)) / sum(L_y_dyn(:,s,t));
        end

        alpha_1 = ave_epsilon_aux.*alpha_0;

        error_0 = max(abs(alpha_0 - alpha_1));

        alpha_0 = 0.8*alpha_0 + 0.2*alpha_1;

        if mod(count,10)==0
            error_0
        end

    end
    
    alpha_nomigr_dyn(:,t) = alpha_0';
    epsilon_nomigr_dyn(:,:,t) = epsilon_aux_0.^(1/theta);
    
end

alpha_hat_nomigr_dyn = log(alpha_nomigr_dyn(:,first_period+1:first_period+T_aux-1)) - log(alpha_nomigr_dyn(:,first_period:first_period+T_aux-2));
